package cn.alm.ragdemo.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class RagConfig {

    @Bean
    public ChatClient chatClient(ChatClient.Builder builder) {
        ChatClient client = builder.defaultSystem("你是一名家庭管家。").build();
        return client;
    }

    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        // 读取文本内容
        String filename = "prompt_template.txt";
        TextReader textReader = new TextReader(filename);
        textReader.getCustomMetadata().put("filename", filename);
        List<Document> documents = textReader.get();
        // 段落切分
        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter(500, 100, 5, 100, true);
        tokenTextSplitter.apply(documents);
        // 将文本添加到向量数据库
        vectorStore.add(documents);
        return vectorStore;
    }

}
